import streamlit as st
import psycopg2
import pandas as pd

# 设置页面标题和图标
st.set_page_config(
    page_title="query_pgsql_schema",
    page_icon="📊",
    layout="wide"
)

# 页面标题
st.title("📊 数据库表结构查询工具")
st.caption("作者：何双新 ｜ 功能：快速查看 PostgreSQL 表结构 | 支持连接检测")
# --- 数据库连接配置区域 ---
with st.sidebar:
    st.header("🔧 数据库连接配置")
    col1, col2 = st.columns(2)
    with col1:
        db_host = st.text_input("主机 (host)", value="localhost")
    with col2:
        db_port = st.text_input("端口 (port)", value="5432")

    db_name = st.text_input("数据库名 (dbname)", value="odoo181121")
    col3, col4 = st.columns(2)
    with col3:
        db_user = st.text_input("用户名 (user)", value="dtcloud")
    with col4:
        db_password = st.text_input("密码 (password)", type="password")

    # 封装连接参数
    db_params = {
        'host': db_host,
        'port': db_port,
        'dbname': db_name,
        'user': db_user,
        'password': db_password
    }

    # 连接检查功能
    st.divider()
    st.subheader("🔍 连接检测")
    if st.button("检查数据库连接"):
        if not all([db_host, db_port, db_name, db_user, db_password]):
            st.error("❗ 请完整填写数据库连接信息")
        else:
            with st.spinner("正在检查连接，请稍候..."):
                try:
                    conn = psycopg2.connect(**db_params)
                    conn.close()
                    st.success("✅ 数据库连接成功！")
                except Exception as e:
                    st.error(f"❌ 数据库连接失败: {e}")

    # 查询表名
    st.divider()
    st.subheader("📋 表名查询")
    table_name = st.text_input("请输入要查询的表名", value="sale_order")

    query_table = st.button("查询表结构")


# --- 查询功能函数 ---
def get_table_structure(db_params, table_name):
    query = """
    SELECT
        a.attname AS column_name,
        pg_catalog.format_type(a.atttypid, a.atttypmod) AS data_type,
        a.attnotnull AS is_not_null,
        a.atthasdef AS has_default,
        pg_get_expr(d.adbin, d.adrelid) AS default_value,
        pg_get_serial_sequence(c.relname, a.attname) IS NOT NULL AS is_serial,
        CASE
            WHEN tc.contype = 'p' THEN 'PRIMARY KEY'
            WHEN tc.contype = 'u' THEN 'UNIQUE'
            WHEN tc.contype = 'f' THEN 'FOREIGN KEY'
            WHEN tc.contype = 'c' THEN 'CHECK'
            ELSE NULL
        END AS constraint_type,
        fk_table.relname AS fk_table,
        fk_column.attname AS fk_column
    FROM
        pg_attribute a
    JOIN pg_class c ON a.attrelid = c.oid
    JOIN pg_namespace n ON c.relnamespace = n.oid
    LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
    LEFT JOIN (
        SELECT
            conrelid,
            contype,
            confrelid,
            conkey,
            confkey
        FROM
            pg_constraint
    ) tc ON c.oid = tc.conrelid AND a.attnum = ANY(tc.conkey)
    LEFT JOIN pg_class fk_table ON tc.confrelid = fk_table.oid
    LEFT JOIN pg_attribute fk_column ON fk_table.oid = fk_column.attrelid AND fk_column.attnum = ANY(tc.confkey)
    WHERE
        c.relname = %s
        AND a.attnum > 0
        AND NOT a.attisdropped
    ORDER BY
        a.attnum;
    """
    try:
        conn = psycopg2.connect(**db_params)
        cursor = conn.cursor()
        cursor.execute(query, (table_name,))
        columns = cursor.fetchall()
        cursor.close()
        conn.close()
        return columns
    except Exception as e:
        st.error(f"数据库查询失败: {e}")
        return None

def format_table_structure(columns):
    """格式化输出表结构"""
    formatted_columns = []
    for col in columns:
        column_name, data_type, is_not_null, has_default, default_value, is_serial, constraint_type, fk_table, fk_column = col
        column_str = f"    {column_name} {data_type}"
        if is_not_null:
            column_str += " NOT NULL"
        if has_default and default_value is not None:
            column_str += f" DEFAULT {default_value}"
        if constraint_type == 'PRIMARY KEY':
            column_str += " PRIMARY KEY"
        if constraint_type == 'UNIQUE':
            column_str += " UNIQUE"
        if constraint_type == 'FOREIGN KEY' and fk_table:
            referenced_column = fk_column if fk_column == 'id' else 'id'
            column_str += f" REFERENCES {fk_table}({referenced_column})"
        formatted_columns.append(column_str)
    return f"(\n" + ",\n".join(formatted_columns) + "\n)"

# --- 查询逻辑处理 ---
if query_table:
    if not all([db_host, db_port, db_name, db_user, db_password]):
        st.error("❗ 请先填写完整数据库连接信息")
    elif not table_name:
        st.error("❗ 请输入表名")
    else:
        with st.spinner("正在查询表结构，请稍候..."):
            columns = get_table_structure(db_params, table_name)

        if columns:
            formatted_structure = format_table_structure(columns)

            st.success(f"✅ 表结构查询成功: `{table_name}`")
            st.code(f"{table_name} {formatted_structure}", language="sql")

            # 表格展示
            st.subheader("🔎 表结构详细信息")
            column_names = [
                "列名", "数据类型", "非空", "有默认值", "默认值", "是否自增", "约束类型", "外键表", "外键列"
            ]
            df = pd.DataFrame(columns, columns=column_names)
            st.dataframe(df, use_container_width=True)
        else:
            st.warning("⚠️ 没有查询到表结构，请检查表名或连接信息。")
st.caption("安徽智加数字科技有限公司 · 技术学习组出品 🚀")
